#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""
简化版ERNIE模型测试程序
"""

import json
import os

def simple_test():
    """
    简单测试ERNIE模型文件是否存在和基本结构
    """
    model_path = "./ERNIE-4.5-0.3B-Paddle"
    
    print("检查ERNIE-4.5-0.3B-Paddle模型文件...")
    
    # 检查模型目录是否存在
    if not os.path.exists(model_path):
        print(f"错误: 模型目录 {model_path} 不存在")
        return
    
    print(f"模型目录 {model_path} 存在")
    
    # 检查关键文件
    key_files = [
        "config.json",
        "tokenizer_config.json",
        "tokenizer.model",
        "model.safetensors",
        "README.md"
    ]
    
    for file in key_files:
        file_path = os.path.join(model_path, file)
        if os.path.exists(file_path):
            size = os.path.getsize(file_path)
            print(f"✓ {file} 存在 (大小: {size//1024} KB)")
        else:
            print(f"✗ {file} 不存在")
    
    # 读取配置文件
    print("\n读取模型配置信息...")
    try:
        with open(os.path.join(model_path, "config.json"), "r", encoding="utf-8") as f:
            config = json.load(f)
            print(f"模型类型: {config.get('model_type', '未知')}")
            print(f"隐藏层大小: {config.get('hidden_size', '未知')}")
            print(f"层数: {config.get('num_hidden_layers', '未知')}")
            print(f"注意力头数: {config.get('num_attention_heads', '未知')}")
    except Exception as e:
        print(f"读取配置文件时出错: {e}")
    
    # 读取tokenizer配置
    print("\n读取Tokenizer配置信息...")
    try:
        with open(os.path.join(model_path, "tokenizer_config.json"), "r", encoding="utf-8") as f:
            tokenizer_config = json.load(f)
            print(f"Tokenizer类: {tokenizer_config.get('tokenizer_class', '未知')}")
            print(f"特殊token:")
            for key, value in tokenizer_config.items():
                if "token" in key.lower() and isinstance(value, str):
                    print(f"  {key}: {value}")
    except Exception as e:
        print(f"读取tokenizer配置文件时出错: {e}")

    print("\n模型检查完成。要运行完整的推理测试，需要安装完整的PaddlePaddle环境。")

if __name__ == "__main__":
    simple_test()